from paddlex import create_pipeline


def paddlex_ocr(image_path,pipeline_yaml="../config_yaml/OCR.yaml",device="gpu",save=False,save_dir="./output/ocr_pipelines/"):
    """
    使用PaddleX的OCR模型进行文本识别
    Args:
        image_path (str): 输入图片路径
        pipeline_yaml (str): OCR模型配置文件路径
        device (str): 设备类型，默认为"gpu"
        save (bool): 是否保存识别结果，默认为False
        save_dir (str): 保存结果的目录，默认为"./output/ocr_pipelines/"
    Returns:
        output: 识别结果
    """
    import os
    print("当前工作目录:", os.getcwd())

    # 创建OCR模型
    pipeline = create_pipeline(pipeline=pipeline_yaml, device=device)

    # 进行OCR识别
    output = pipeline.predict(
        input=image_path,
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
    )
    output = list(output)
    if save:
        for res in output:
            res.save_to_img(save_dir)
            res.save_to_json(save_dir)
    return output

if __name__ == "__main__":
    # 示例用法
    image_path = ["/home/fengjie/doc-parser/MinerU/src/pymupdf_text/提取自德州德达城市建设投资运营有限公司2023年面向专业投资者公开发行公司债券（第一期）募集说明书.pdf"]  # 输入图片路径列表
    output = paddlex_ocr(image_path,save=True)
    for res in output:
        res.print()
